#include <rclcpp/rclcpp.hpp>
#include <livox_interfaces/msg/custom_msg.hpp>
#include <sensor_msgs/msg/point_cloud2.hpp>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/filters/passthrough.h>
#include <pcl/filters/extract_indices.h>
#include <pcl/segmentation/extract_clusters.h>
#include <pcl/kdtree/kdtree.h>
#include <pcl/filters/voxel_grid.h>
#include <cmath>
#include <thread>
#include <mutex>
#include <chrono>
#include <deque>
#include <atomic>

// Open3D可视化头文件
#include <open3d/Open3D.h>

class LidarProcessor : public rclcpp::Node
{
public:
    LidarProcessor() : Node("lidar_processor"), 
                       is_running_(true),
                       frame_count_(0)
    {
        // 声明并获取ROS参数
        this->declare_parameter("input_cloud", "/livox/lidar_avia");
        this->declare_parameter("distance_threshold", 30.0);
        this->declare_parameter("angle_threshold", 15.0);
        this->declare_parameter("voxel_size", 0.3);
        this->declare_parameter("cluster_eps", 1.0);
        this->declare_parameter("cluster_min_points", 10);
        this->declare_parameter("enable_clustering", false);  // 默认关闭聚类提高速度

        // 获取参数
        std::string input_topic = this->get_parameter("input_cloud").as_string();
        distance_threshold_ = this->get_parameter("distance_threshold").as_double();
        angle_threshold_ = this->get_parameter("angle_threshold").as_double();
        voxel_size_ = this->get_parameter("voxel_size").as_double();
        cluster_eps_ = this->get_parameter("cluster_eps").as_double();
        cluster_min_points_ = this->get_parameter("cluster_min_points").as_int();
        enable_clustering_ = this->get_parameter("enable_clustering").as_bool();

        exclude_front_ = true;
        cos_angle_threshold_ = std::cos(angle_threshold_ * M_PI / 180.0);

        // 创建订阅者
        sub_ = this->create_subscription<livox_interfaces::msg::CustomMsg>(
            input_topic, 50, std::bind(&LidarProcessor::callback, this, std::placeholders::_1));

        // 创建发布者
        pcl_pub_ = this->create_publisher<sensor_msgs::msg::PointCloud2>("/livox/pointcloud2", 10);

        RCLCPP_INFO(this->get_logger(), "订阅: %s，发布可视化点云到: /livox/pointcloud2", input_topic.c_str());
        RCLCPP_INFO(this->get_logger(), "参数设置: 距离阈值=%.1f, 体素大小=%.2f, 聚类=%s",
                   distance_threshold_, voxel_size_, enable_clustering_ ? "开启" : "关闭");

        // 帧率计算参数
        fps_window_size_ = 10;
        last_fps_print_time_ = this->now();

        // 初始化点云共享指针 - 创建一个初始点以便窗口不黑屏
        dummy_cloud_ = std::make_shared<open3d::geometry::PointCloud>();
        dummy_cloud_->points_.push_back(Eigen::Vector3d(0, 0, 0));  // 原点
        dummy_cloud_->colors_.push_back(Eigen::Vector3d(1, 1, 1));  // 白色
        
        // 创建处理线程和可视化线程，分离数据处理和显示
        processing_thread_ = std::thread(&LidarProcessor::processing_loop, this);
        visualization_thread_ = std::thread(&LidarProcessor::visualization_loop, this);
    }

    ~LidarProcessor()
    {
        is_running_ = false;
        if (processing_thread_.joinable()) {
            processing_thread_.join();
        }
        if (visualization_thread_.joinable()) {
            visualization_thread_.join();
        }
    }

private:
    // 回调函数只负责接收数据，不做处理，避免阻塞
    void callback(const livox_interfaces::msg::CustomMsg::SharedPtr msg)
    {
        std::lock_guard<std::mutex> lock(queue_mutex_);
        // 如果队列太大，丢弃旧数据
        if (msg_queue_.size() > 10) {
            msg_queue_.pop_front();
            RCLCPP_WARN(this->get_logger(), "消息队列堆积，丢弃旧数据");
        }
        msg_queue_.push_back(*msg);  // 存储拷贝，避免共享指针问题
    }

    // 独立的处理线程，负责点云处理
    void processing_loop()
    {
        while (is_running_ && rclcpp::ok()) {
            // 检查是否有新消息
            if (msg_queue_.empty()) {
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
                continue;
            }

            auto frame_start_time = std::chrono::high_resolution_clock::now();
            livox_interfaces::msg::CustomMsg msg;
            
            // 取出消息
            {
                std::lock_guard<std::mutex> lock(queue_mutex_);
                msg = msg_queue_.front();
                msg_queue_.pop_front();
            }

            // 转换为PCL点云
            pcl::PointCloud<pcl::PointXYZI>::Ptr cloud(new pcl::PointCloud<pcl::PointXYZI>);
            cloud->points.reserve(msg.points.size());
            for (const auto& p : msg.points) {
                pcl::PointXYZI point;
                point.x = p.x;
                point.y = p.y;
                point.z = p.z;
                point.intensity = p.reflectivity;
                cloud->points.push_back(point);
            }

            frame_count_++;
            if (frame_count_ % 10 == 0) {  // 每10帧打印一次点云数量
                RCLCPP_INFO(this->get_logger(), "原始点云数量: %zu", cloud->points.size());
            }

            // 如果没有点，跳过处理
            if (cloud->points.empty()) {
                continue;
            }

            // 1. 快速距离过滤
            pcl::PointCloud<pcl::PointXYZI>::Ptr filtered_cloud(new pcl::PointCloud<pcl::PointXYZI>);
            pcl::PassThrough<pcl::PointXYZI> pass;
            pass.setInputCloud(cloud);
            pass.setFilterFieldName("x");
            pass.setFilterLimits(-distance_threshold_, distance_threshold_);
            pass.filter(*filtered_cloud);

            // 2. Z轴过滤
            pass.setInputCloud(filtered_cloud);
            pass.setFilterFieldName("z");
            pass.setFilterLimits(-distance_threshold_/2, distance_threshold_/2);
            pass.filter(*filtered_cloud);

            // 3. 方向过滤
            if (exclude_front_) {
                pcl::PointIndices::Ptr inliers(new pcl::PointIndices);
                inliers->indices.reserve(filtered_cloud->size()/2);
                for (size_t i = 0; i < filtered_cloud->points.size(); ++i) {
                    const auto& point = filtered_cloud->points[i];
                    double dist_sq = point.x * point.x + point.y * point.y + point.z * point.z;
                    if (dist_sq > 0.1) {
                        double cos_angle = point.x / std::sqrt(dist_sq);
                        if (std::abs(cos_angle) < cos_angle_threshold_) {
                            inliers->indices.push_back(i);
                        }
                    }
                }
                pcl::ExtractIndices<pcl::PointXYZI> extract;
                extract.setInputCloud(filtered_cloud);
                extract.setIndices(inliers);
                extract.filter(*filtered_cloud);
            }

            // 4. 降采样
            pcl::PointCloud<pcl::PointXYZI>::Ptr downsampled_cloud(new pcl::PointCloud<pcl::PointXYZI>);
            if (filtered_cloud->points.size() > 100) {
                pcl::VoxelGrid<pcl::PointXYZI> sor;
                sor.setInputCloud(filtered_cloud);
                sor.setLeafSize(voxel_size_, voxel_size_, voxel_size_);
                sor.filter(*downsampled_cloud);
            } else {
                *downsampled_cloud = *filtered_cloud;
            }

            // 5. 聚类（可选）
            if (enable_clustering_ && downsampled_cloud->points.size() > 100) {
                pcl::search::KdTree<pcl::PointXYZI>::Ptr tree(new pcl::search::KdTree<pcl::PointXYZI>);
                tree->setInputCloud(downsampled_cloud);

                std::vector<pcl::PointIndices> cluster_indices;
                pcl::EuclideanClusterExtraction<pcl::PointXYZI> ec;
                ec.setClusterTolerance(cluster_eps_);
                ec.setMinClusterSize(cluster_min_points_);
                ec.setMaxClusterSize(10000);
                ec.setSearchMethod(tree);
                ec.setInputCloud(downsampled_cloud);
                ec.extract(cluster_indices);
            }

            // 发布处理后的点云
            sensor_msgs::msg::PointCloud2 pcl_msg;
            pcl::toROSMsg(*filtered_cloud, pcl_msg);
            pcl_msg.header = msg.header;
            pcl_pub_->publish(pcl_msg);

            // 将处理后的点云传递给可视化线程
            {
                std::lock_guard<std::mutex> lock(visualization_mutex_);
                // 转换为Open3D点云
                auto o3d_cloud = std::make_shared<open3d::geometry::PointCloud>();
                o3d_cloud->points_.resize(downsampled_cloud->size());
                o3d_cloud->colors_.reserve(downsampled_cloud->size());
                
                double min_intensity = 0.0, max_intensity = 255.0;
                for (size_t i = 0; i < downsampled_cloud->size(); ++i) {
                    const auto& p = downsampled_cloud->points[i];
                    o3d_cloud->points_[i] = Eigen::Vector3d(p.x, p.y, p.z);
                    
                    double norm_intensity = std::min(1.0, std::max(0.0, (p.intensity - min_intensity) / (max_intensity - min_intensity)));
                    o3d_cloud->colors_.push_back(Eigen::Vector3d(norm_intensity, 0.5*(1-norm_intensity), 1-norm_intensity));
                }
                
                visualization_cloud_ = o3d_cloud;
                new_cloud_available_ = true;
            }

            // 计算帧率
            calculate_fps(frame_start_time);
        }
    }

    // 独立的可视化线程
    void visualization_loop()
    {
        open3d::visualization::Visualizer vis;
        if (!vis.CreateVisualizerWindow("LiDAR Point Cloud Visualization", 1024, 768)) {
            RCLCPP_ERROR(this->get_logger(), "无法创建可视化窗口！");
            return;
        }
        
        // 添加初始点云
        vis.AddGeometry(dummy_cloud_);
        
        // 优化渲染设置
        auto& render_option = vis.GetRenderOption();
        render_option.background_color_ = Eigen::Vector3d(0, 0, 0);  // 黑色背景
        render_option.point_size_ = 2.0;
        
        bool first_update = true;
        std::shared_ptr<open3d::geometry::PointCloud> current_cloud = dummy_cloud_;
        
        while (is_running_ && rclcpp::ok()) {
            // 检查是否有新的点云数据
            {
                std::lock_guard<std::mutex> lock(visualization_mutex_);
                if (new_cloud_available_) {
                    // 移除旧点云
                    vis.RemoveGeometry(current_cloud);
                    
                    // 更新为新点云
                    current_cloud = visualization_cloud_;
                    vis.AddGeometry(current_cloud);
                    
                    // 第一次更新时调整视角
                    if (first_update) {
                        auto& view_control = vis.GetViewControl();
                        view_control.SetFront(Eigen::Vector3d(0, -1, 0));
                        view_control.SetLookat(Eigen::Vector3d(0, 0, -1));
                        view_control.SetUp(Eigen::Vector3d(0, 0, 1));
                        view_control.SetZoom(0.3);
                        first_update = false;
                    }
                    
                    new_cloud_available_ = false;
                    
                    // 打印当前点云数量，帮助调试
                    RCLCPP_INFO(this->get_logger(), "可视化点云数量: %zu", current_cloud->points_.size());
                }
            }
            
            // 刷新窗口
            vis.PollEvents();
            vis.UpdateRender();
            
            // 控制可视化帧率
            std::this_thread::sleep_for(std::chrono::milliseconds(33));  // ~30 FPS
        }
        
        vis.DestroyVisualizerWindow();
    }

    // 计算帧率
    void calculate_fps(const std::chrono::high_resolution_clock::time_point& frame_start)
    {
        auto frame_end = std::chrono::high_resolution_clock::now();
        double frame_time_ms = std::chrono::duration<double, std::milli>(frame_end - frame_start).count();
        
        {
            std::lock_guard<std::mutex> lock(fps_mutex_);
            frame_times_.push_back(frame_time_ms);
            if (frame_times_.size() > fps_window_size_) {
                frame_times_.pop_front();
            }
        }
        
        auto now = this->now();
        auto elapsed = now - last_fps_print_time_;
        if (elapsed.seconds() >= 1.0) {
            std::lock_guard<std::mutex> lock(fps_mutex_);
            
            double total_time = 0;
            for (double t : frame_times_) {
                total_time += t;
            }
            
            double avg_fps = 0;
            if (total_time > 0) {
                avg_fps = frame_times_.size() / (total_time / 1000.0);
            }
            
            RCLCPP_INFO(this->get_logger(), "平均帧率: %.1f FPS, 平均处理时间: %.1f ms", 
                       avg_fps, total_time / frame_times_.size());
            
            last_fps_print_time_ = now;
        }
    }

    // ROS相关成员
    rclcpp::Subscription<livox_interfaces::msg::CustomMsg>::SharedPtr sub_;
    rclcpp::Publisher<sensor_msgs::msg::PointCloud2>::SharedPtr pcl_pub_;

    // 数据处理参数
    double distance_threshold_;
    double angle_threshold_;
    bool exclude_front_;
    double cos_angle_threshold_;
    double cluster_eps_;
    int cluster_min_points_;
    bool enable_clustering_;
    double voxel_size_;

    // 多线程相关
    std::thread processing_thread_;    // 点云处理线程
    std::thread visualization_thread_; // 可视化线程
    std::atomic<bool> is_running_;     // 线程运行标志
    size_t frame_count_;               // 帧计数器

    // 消息队列 - 用于接收线程和处理线程之间的通信
    std::deque<livox_interfaces::msg::CustomMsg> msg_queue_;
    std::mutex queue_mutex_;

    // 可视化相关
    std::mutex visualization_mutex_;
    std::shared_ptr<open3d::geometry::PointCloud> visualization_cloud_;
    std::shared_ptr<open3d::geometry::PointCloud> dummy_cloud_;
    std::atomic<bool> new_cloud_available_;

    // 帧率计算
    std::mutex fps_mutex_;
    std::deque<double> frame_times_;
    size_t fps_window_size_;
    rclcpp::Time last_fps_print_time_;
};

int main(int argc, char * argv[])
{
    rclcpp::init(argc, argv);
    auto node = std::make_shared<LidarProcessor>();
    rclcpp::spin(node);
    rclcpp::shutdown();
    return 0;
}
    